package com.yanx.common.utils;

import java.util.*;

public final class ThreadLocal_ {
    private static final ThreadLocal<Map<String, Object>> threadLocal = ThreadLocal.withInitial(() -> new HashMap(4));

    public static Map<String, Object> getThreadLocal() {
        return threadLocal.get();
    }

    public static <T> T get(String key) {
        return (T) getThreadLocal().get(key);
    }

    public static <T> T get(String key, Class<T> t) {
        return (T) get(key);
    }

    public static <T> T get(String key, T defaultValue) {
        T obj = get(key);
        return obj == null ? defaultValue : obj;
    }

    public static <T> Map<String, T> getByPrefix(String prefix) {
        Map<String, T> vars = new HashMap<>();
        if (prefix == null) {
            return vars;
        }
        Set<Map.Entry<String, Object>> set = getThreadLocal().entrySet();

        for (Map.Entry<String, Object> entry : set) {
            Object key = entry.getKey();
            if (key instanceof String) {
                String strKey = (String) key;
                if (strKey.startsWith(prefix)) {
                    vars.put(strKey, (T) entry.getValue());
                }
            }
        }
        return vars;
    }

    public static <T> T getAndRemove(String key) {
        T res = get(key);
        remove(key);
        return res;
    }

    public static <T> T getAndRemove(String key, Class<T> t) {
        T res = get(key, t);
        remove(key);
        return res;
    }

    public static <T> T getAndRemove(String key, T defaultValue) {
        T res = get(key);
        remove(key);
        return res == null ? defaultValue : res;
    }

    public static <T> Map<String, T> getAndRemoveByPrefix(String prefix) {
        Map<String, T> map = getByPrefix(prefix);
        removeByPrefix(prefix);
        return map;
    }

    public static void set(String key, Object value) {
        getThreadLocal().put(key, value);
    }

    public static void set(Map<String, Object> keyValueMap) {
        getThreadLocal().putAll(keyValueMap);
    }

    public static void remove() {
        threadLocal.remove();
    }

    public static <T> T remove(String key) {
        return (T) getThreadLocal().remove(key);
    }

    public static void removeByPrefix(String prefix) {
        if (prefix == null) {
            return;
        }
        Set<Map.Entry<String, Object>> set = getThreadLocal().entrySet();
        List<String> removeKeys = new ArrayList<>();

        for (Map.Entry<String, Object> entry : set) {
            Object key = entry.getKey();
            if (key instanceof String) {
                String strKey = (String) key;
                if (strKey.startsWith(prefix)) {
                    removeKeys.add(strKey);
                }
            }
        }
        for (String key : removeKeys) {
            remove(key);
        }
    }
}